package ao.ai.demo.supervised.classification.plane;

import ao.ai.classify.decision.MinInfoBinClassifier;
import ao.ai.model.common.data.impl.IntNumList;
import ao.ai.model.ml.classify.example.impl.XyBinClassified;
import ao.ai.model.ml.classify.hypothesis.classifier.num.NumBinClassifier;
import ao.ai.model.ml.classify.learner.num.NumBinProbClassLeaner;
import ao.util.async.Condition;
import com.google.common.collect.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.swing.*;
import java.awt.*;
import java.awt.event.MouseEvent;
import java.awt.event.MouseListener;
import java.awt.event.MouseMotionListener;
import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.List;

/**
 * User: aostrovsky
 * Date: 5-Feb-2010
 * Time: 11:45:22 PM
 */
public class BinaryClassificationPanel
        extends    JPanel
        implements MouseListener,
                   MouseMotionListener
{
    //-------------------------------------------------------------------------
    private static final Logger LOG =
            LoggerFactory.getLogger( BinaryClassificationPanel.class );

    private static final Color BACK_GROUND    = Color.BLACK;
    private static final Color POSITIVE       = Color.RED;
    private static final Color POSITIVE_GUESS = POSITIVE.darker();
    private static final Color NEGATIVE       = Color.BLUE;
    private static final Color NEGATIVE_GUESS = NEGATIVE.darker();


    //-------------------------------------------------------------------------
    public static void main(String[] args)
            throws InvocationTargetException, InterruptedException
    {
        final BinaryClassificationPanel classPanel =
                new BinaryClassificationPanel();

        JFrame frame = new JFrame();
        frame.setDefaultCloseOperation( JFrame.EXIT_ON_CLOSE );
        frame.setContentPane( classPanel );
        frame.pack();
        frame.setSize(800, 500);
        frame.setVisible( true );

        for (long i = 0; i < Long.MAX_VALUE; i++)
        {
            classPanel.waitForNewExamples();
            classPanel.makePredictions();
            classPanel.refresh(false);

//            SwingUtilities.invokeAndWait(new Runnable() {
////            SwingUtilities.invokeLater(new Runnable() {
//                @Override public void run() {
//                    classPanel.repaint();
//                }
//            });
        }
    }


    //-------------------------------------------------------------------------
    private final List<XyBinClassified> examples;
    private final List<XyBinClassified> predictions;

    private final Condition available;


    //-------------------------------------------------------------------------
    public BinaryClassificationPanel()
    {
        examples    = new ArrayList<XyBinClassified>();
        predictions = new ArrayList<XyBinClassified>();

        available   = new Condition(false);

        addMouseListener( this );
        addMouseMotionListener( this );
    }


    //-------------------------------------------------------------------------
    public synchronized void clear()
    {
        examples.clear();
        predictions.clear();
    }


    //-------------------------------------------------------------------------
    private void waitForNewExamples()
    {
        try {
            available.waitForTrue();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        available.setFalse();
    }


    //-------------------------------------------------------------------------
    private void makePredictions()
    {
        final ArrayList<XyBinClassified> copyOfExamples;
        synchronized (this)
        {
            copyOfExamples = Lists.newArrayList(examples);
        }

        if (copyOfExamples.size() < 1) {
            return;
        }

        NumBinProbClassLeaner learner =
                new MinInfoBinClassifier();

//        SupervisedPlainLearner<
//                    NumericalFeatureList,
//                    SingleBinaryFeature>
//                learner =
//                    new MeanBinaryClassifier();
//                    new LogisticStochasticGradientAscent();
//                    new RandomForestBinaryClassifier();
//                    new MinimumInformationBinaryClassifier();

        NumBinClassifier hype =
                learner.learn(
                        copyOfExamples);

//        SupervisedHypothesis<
//                NumericalFeatureList,
//                SingleBinaryFeature>
//            hype = learner.learn(
//                copyOfExamples);

        final List<XyBinClassified> nextPrediction =
                new ArrayList<XyBinClassified>();

        for (int w = 0; w < getWidth(); w += 10)
        {
            for (int h = 0; h < getHeight(); h += 10)
            {
                try {
                    boolean isPositive =
//                            hype.predict(
//                                    new FeatureVector(null, w, h)
//                            ).binaryCategory();
                            hype.classify(
                                    new IntNumList(w, h)
                            ).best().getBin();

                    nextPrediction.add(
                            new XyBinClassified(
                                    w, h, isPositive));
                } catch (Throwable ignored) {}
            }
        }

        synchronized (this)
        {
            predictions.clear();
            predictions.addAll( nextPrediction );
        }
    }


    //-------------------------------------------------------------------------
    @Override
    public synchronized void paint(Graphics graphics)
    {
        graphics.setColor( BACK_GROUND );
        graphics.fillRect(
                0, 0, getWidth(), getHeight());

        paint(graphics, POSITIVE      , NEGATIVE      , examples   , 2);
        paint(graphics, POSITIVE_GUESS, NEGATIVE_GUESS, predictions, 1);
    }

    private void paint(
            Graphics              graphics,
            Color                 positiveColour,
            Color                 negativeColour,
            List<XyBinClassified> examples,
            int                   size)
    {
        for (XyBinClassified point : examples)
        {
            graphics.setColor(
                    point.classification().best().getBin()
                    ? positiveColour : negativeColour);

            double[] input = point.input().getDoubles();
            graphics.drawRect(
                    (int) Math.round(input[0]),
                    (int) Math.round(input[1]),
                    size, size);
        }
    }


    //-------------------------------------------------------------------------
    @Override public void mouseReleased(MouseEvent e) {}
    @Override public void mouseEntered (MouseEvent e) {}
    @Override public void mouseExited  (MouseEvent e) {}
    @Override public void mouseMoved   (MouseEvent e) {}

    @Override public void mouseClicked (MouseEvent e)
    {
        if (e.getClickCount() == 3)
        {
            clear();
        }
    }

    private boolean previousClicked = false;

    @Override public void mousePressed (MouseEvent e) {
        previousClicked = isPositive(e);
        userInput( e );
    }

    @Override public void mouseDragged(MouseEvent e) {
        userInput( e );
    }

    private boolean isPositive(MouseEvent e)
    {
        return (e.getButton() == MouseEvent.BUTTON1);
    }

    private synchronized void userInput(MouseEvent e)
    {
        examples.add(new XyBinClassified(
                e.getX(), e.getY(), previousClicked));

        available.setTrue();
        refresh(true);
    }


    //-------------------------------------------------------------------------
    private void refresh(final boolean onlyIfAvailable)
    {
        SwingUtilities.invokeLater(new Runnable() {
            @Override public void run() {
                if (onlyIfAvailable &&
                    ! available.isTrue()) return;
                repaint();
            }
        });
    }
}